import numpy as np
from scipy.sparse import lil_matrix, csr_matrix
from scipy.sparse.linalg import spsolve
import matplotlib.pyplot as plt
import scienceplots
plt.style.use(['science'])

class FEM2D_Poisson:
    def __init__(self, Lx=1, Ly=1, nx=20, ny=20):
        self.Lx, self.Ly = Lx, Ly
        self.nx, self.ny = nx, ny
        self.hx = Lx / nx
        self.hy = Ly / ny
        self.nodes, self.elements = self.generate_mesh()

    def generate_mesh(self):
        x = np.linspace(0, self.Lx, self.nx + 1)
        y = np.linspace(0, self.Ly, self.ny + 1)
        nodes = np.array(np.meshgrid(x, y)).T.reshape(-1, 2)

        elements = []
        for i in range(self.ny):
            for j in range(self.nx):
                n0 = i * (self.nx + 1) + j
                n1 = n0 + 1
                n2 = (i + 1) * (self.nx + 1) + j + 1
                n3 = (i + 1) * (self.nx + 1) + j
                elements.append([n0, n1, n2, n3])

        return nodes, np.array(elements)

    def source_term(self, x, y):
        return -6 * x * y

    def assemble_system(self):
        n_nodes = len(self.nodes)
        K = lil_matrix((n_nodes, n_nodes))
        F = np.zeros(n_nodes)

        gauss_points = np.array([[-1, -1], [1, -1], [1, 1], [-1, 1]]) * 0.57735
        weights = [1, 1, 1, 1]

        for elem in self.elements:
            node_coords = self.nodes[elem]
            Ke = np.zeros((4, 4))
            Fe = np.zeros(4)

            for gp, w in zip(gauss_points, weights):
                xi, eta = gp
                N = 0.25 * np.array([(1 - xi) * (1 - eta),
                                     (1 + xi) * (1 - eta),
                                     (1 + xi) * (1 + eta),
                                     (1 - xi) * (1 + eta)])
                dN_dxi = 0.25 * np.array([-(1 - eta), (1 - eta), (1 + eta), -(1 + eta)])
                dN_deta = 0.25 * np.array([-(1 - xi), -(1 + xi), (1 + xi), (1 - xi)])

                J = np.dot(np.array([dN_dxi, dN_deta]), node_coords)
                detJ = np.linalg.det(J)
                invJ = np.linalg.inv(J)

                dN_dx = invJ[0, 0] * dN_dxi + invJ[0, 1] * dN_deta
                dN_dy = invJ[1, 0] * dN_dxi + invJ[1, 1] * dN_deta

                Ke += (np.outer(dN_dx, dN_dx) + np.outer(dN_dy, dN_dy)) * detJ * w

                x_phys = np.dot(N, node_coords[:, 0])
                y_phys = np.dot(N, node_coords[:, 1])
                Fe += N * -self.source_term(x_phys, y_phys) * detJ * w

            for i in range(4):
                for j in range(4):
                    K[elem[i], elem[j]] += Ke[i, j]
                F[elem[i]] += Fe[i]

        return K.tocsr(), F

    def apply_dirichlet_bc(self, K, F):
        tol = 1e-6
        boundary_nodes = []
        boundary_nodes += np.where(self.nodes[:, 0] < tol)[0].tolist()
        boundary_nodes += np.where(self.nodes[:, 0] > self.Lx - tol)[0].tolist()
        boundary_nodes += np.where(self.nodes[:, 1] < tol)[0].tolist()
        boundary_nodes += np.where(self.nodes[:, 1] > self.Ly - tol)[0].tolist()

        for node in boundary_nodes:
            K[node, :] = 0
            K[:, node] = 0
            K[node, node] = 1
            F[node] = 0

        return K, F

    def solve(self):
        K, F = self.assemble_system()
        K, F = self.apply_dirichlet_bc(K, F)
        u = spsolve(K, F)
        return u

    def exact_solution(self, n_max=20):
        """计算精确解"""
        x = self.nodes[:, 0]
        y = self.nodes[:, 1]
        u_exact = np.zeros_like(x)

        # 多项式部分
        polynomial = (1 - x ** 2) * x * y

        # 级数部分
        series_sum = np.zeros_like(x)
        for n in range(1, n_max + 1):
            numerator = (-1) ** n * 12 * np.sinh(n * np.pi * y)
            denominator = n ** 3 * np.pi ** 3 * np.sinh(n * np.pi)
            sin_term = np.sin(n * np.pi * x)
            series_sum += (numerator / denominator) * sin_term

        u_exact = polynomial + series_sum
        return u_exact

    def compute_errors(self, u_numerical):
        """计算数值解与精确解的误差"""
        u_exact = self.exact_solution()
        error = u_numerical - u_exact
        l2_error = np.sqrt(np.mean(error ** 2))
        max_error = np.max(np.abs(error))
        return l2_error, max_error, error

    def visualize_numerical_solution(self, u_numerical):
        """绘制数值解图像（2D伪彩色图）"""
        X = self.nodes[:, 0].reshape(self.ny + 1, self.nx + 1)
        Y = self.nodes[:, 1].reshape(self.ny + 1, self.nx + 1)
        U_num = u_numerical.reshape(self.ny + 1, self.nx + 1)

        fig = plt.figure(figsize=(4, 3))
        ax = fig.add_subplot(111)
        mesh = ax.pcolormesh(X, Y, U_num, cmap='viridis', shading='auto')
        fig.colorbar(mesh, ax=ax, label='$u(x, y)$')
        ax.set_title('Numerical Solution')
        ax.set_xlabel('$x$')
        ax.set_ylabel('$y$')
        plt.tight_layout()
        plt.show()
        fig.savefig('./numerical_solution2.png', dpi=300)

    def visualize_exact_solution(self):
        """绘制精确解图像（2D伪彩色图）"""
        u_exact = self.exact_solution()
        X = self.nodes[:, 0].reshape(self.ny + 1, self.nx + 1)
        Y = self.nodes[:, 1].reshape(self.ny + 1, self.nx + 1)
        U_exact = u_exact.reshape(self.ny + 1, self.nx + 1)

        fig = plt.figure(figsize=(4, 3))
        ax = fig.add_subplot(111)
        mesh = ax.pcolormesh(X, Y, U_exact, cmap='viridis', shading='auto')
        fig.colorbar(mesh, ax=ax, label='$u(x, y)$')
        ax.set_title('Exact Solution')
        ax.set_xlabel('$x$')
        ax.set_ylabel('$y$')
        plt.tight_layout()
        plt.show()
        fig.savefig('./exact_solution2.png', dpi=300)

    def visualize_error(self, u_numerical):
        """绘制误差分布图像（2D伪彩色图）"""
        _, _, error = self.compute_errors(u_numerical)
        Error = error.reshape(self.ny + 1, self.nx + 1)
        X = self.nodes[:, 0].reshape(self.ny + 1, self.nx + 1)
        Y = self.nodes[:, 1].reshape(self.ny + 1, self.nx + 1)

        fig = plt.figure(figsize=(4, 3))
        ax = fig.add_subplot(111)
        mesh = ax.pcolormesh(X, Y, Error, cmap='coolwarm', shading='auto')
        fig.colorbar(mesh, ax=ax, label='Error')
        # ax.set_title('Error')
        ax.set_xlabel('$x$')
        ax.set_ylabel('$y$')
        plt.tight_layout()
        plt.show()
        fig.savefig('./error2', dpi=300)


if __name__ == "__main__":
    # 创建FEM求解器
    fem = FEM2D_Poisson(nx=30, ny=30)

    # 求解数值解
    u_numerical = fem.solve()

    # 计算并打印误差
    l2_error, max_error, _ = fem.compute_errors(u_numerical)
    print(f"L2 Error: {l2_error:.4e}")
    print(f"Max Error: {max_error:.4e}")

    # 可视化结果
    fem.visualize_numerical_solution(u_numerical)  # 数值解图像
    fem.visualize_exact_solution()  # 精确解图像
    fem.visualize_error(u_numerical)  # 误差分布图像